#include "VideoChannel.h"

#include <rtmp.h>

#include <string.h>
#include <fstream>

ofstream fs("sdcard/1.h264", ios::app);

VideoChannel::VideoChannel(SafeQueue<RTMPPacket *> *&packetQueue) : packetQueue(packetQueue) {
    pthread_mutex_init(&mutex, nullptr);
}

VideoChannel::~VideoChannel() {
    pthread_mutex_destroy(&mutex);
    //释放编码器
    if (codec) {
        x264_encoder_close(codec);
        codec = nullptr;
    }
    if (i420Pic) {
        x264_picture_clean(i420Pic);
        DELETE(i420Pic)
    }
}

//设置视频编码信息
void VideoChannel::setVideoEncodeInfo(int width, int height, int fps, int bitrate) {
    pthread_mutex_lock(&mutex);
    //释放旧的编码器
    if (codec) {
        x264_encoder_close(codec);
        codec = nullptr;
    }
    if (i420Pic) {
        x264_picture_clean(i420Pic);
        DELETE(i420Pic)
    }
    //设置新参数
    this->width = width;
    this->height = height;
    this->fps = fps;
    this->bitrate = bitrate;
    this->ySize = width * height;
    this->uvSize = ySize / 4;
    //创建新的编码器
    x264_param_t param;
    x264_param_default_preset(&param, "ultrafast", "zerolatency"); //最快，无延时
    param.i_level_idc = 32;
    param.i_csp = X264_CSP_I420; //输入视频的采样格式，x264只支持I420，因此需要将采集到的NV21转为I420再输入
    param.i_width = width;
    param.i_height = height;
    param.i_bframe = 0; //无B帧
    param.rc.i_rc_method = X264_RC_ABR; //码率控制策略：平均码率
    param.rc.i_bitrate = bitrate / 1000; //比特率，kbps
    param.rc.i_vbv_max_bitrate = bitrate / 1000 * 1.2; //最大码率
    param.rc.i_vbv_buffer_size = bitrate / 1000; //码率控制区大小
    param.i_fps_num = fps; //帧率
    param.i_fps_den = 1; //时间基
    param.i_timebase_den = param.i_fps_num;
    param.i_timebase_num = param.i_fps_den;
    param.b_vfr_input = 0; //0使用fps来控制速率，1使用时间戳来控制速率
    param.i_keyint_max = fps * 2; //每2s一个关键帧
    param.b_repeat_headers = 1; //每个关键帧都携带sps和pps
    param.i_threads = 1; //开启多线程
    x264_param_apply_profile(&param, "baseline");
    //打开编码器
    codec = x264_encoder_open(&param);
    i420Pic = new x264_picture_t;
    x264_picture_alloc(i420Pic, X264_CSP_I420, width, height);
    pthread_mutex_unlock(&mutex);
}

//编码NV21帧数据为RTMPPacket
void VideoChannel::encodeNV21ToRtmp(int8_t *nv21) {
    pthread_mutex_lock(&mutex);
    //NV21转I420图像
    memcpy(i420Pic->img.plane[0], nv21, ySize); //拷贝Y
    for (int i = 0; i < uvSize; ++i) {
        *(i420Pic->img.plane[1] + i) = *(nv21 + ySize + i * 2 + 1); //拷贝U
        *(i420Pic->img.plane[2] + i) = *(nv21 + ySize + i * 2); //拷贝V
    }
    //I420图像编码为H264帧
    x264_nal_t *naluBuffer; //NALU数据
    int naluSize; //NALU数量
    x264_picture_t h264Pic;
    x264_encoder_encode(codec, &naluBuffer, &naluSize, i420Pic, &h264Pic); //编码I420为H264
    //推送帧数据
    int sps_len;
    int pps_len;
    uint8_t sps[100];
    uint8_t pps[100];
    //RTMP视频数据包，是由一个存储VideoDecodeConfig的FlvTag，和若干存储视频裸数据FlvTag组成的
    //VideoDecodeConfig中存放的是视频参数信息，包括SPS和PPS
    for (int i = 0; i < naluSize; i++) {
        //H264裸数据写入文件
        fs.write((char *) naluBuffer[i].p_payload, naluBuffer[i].i_payload);
        int type = naluBuffer[i].i_type;
        if (type == NAL_SPS) {
            sps_len = naluBuffer[i].i_payload - 4; //排除H264头部的00 00 00 01四个固定字节
            memcpy(sps, naluBuffer[i].p_payload + 4, sps_len); //PPS和SPS肯定是一起的，等待PPS一起发送
        } else if (type == NAL_PPS) {
            pps_len = naluBuffer[i].i_payload - 4;
            memcpy(pps, naluBuffer[i].p_payload + 4, pps_len);
            //创建存储VideoDecodeConfig的Packet
            RTMPPacket *packet = createConfigPacket(sps, pps, sps_len, pps_len);
            packetQueue->enqueue(packet);
        } else {
            //创建存储H264视频裸数据的Packet
            RTMPPacket *packet = createRawDataPacket(naluBuffer[i].i_type, naluBuffer[i].p_payload, naluBuffer[i].i_payload);
            packetQueue->enqueue(packet);
        }
    }
    pthread_mutex_unlock(&mutex);
}

//创建存储VideoDecodeConfig的Packet
RTMPPacket *VideoChannel::createConfigPacket(uint8_t *sps, uint8_t *pps, int sps_len, int pps_len) {
    //创建Packet
    int bodySize = 13 + sps_len + 3 + pps_len;
    RTMPPacket *packet = new RTMPPacket;
    RTMPPacket_Alloc(packet, bodySize);
    //按字节向Packet写入数据
    //RTMPPacket的格式和FLV的格式基本是一样的，只是略有不同
    int i = 0;
    //01表示完整帧（IDR），02表示切片帧（SLICE），07表示x264编码，通过位运算得到0x17
    packet->m_body[i++] = 0x17;
    //00表示AVC Header，01表示AVC NALU
    packet->m_body[i++] = 0x00;
    //时间戳
    packet->m_body[i++] = 0x00;
    packet->m_body[i++] = 0x00;
    packet->m_body[i++] = 0x00;
    //协议版本号
    packet->m_body[i++] = 0x01;
    //从SPS中拿到编码规格
    packet->m_body[i++] = sps[1];
    packet->m_body[i++] = sps[2];
    packet->m_body[i++] = sps[3];
    //NALU长度为03，通过位运算得到0xFF
    packet->m_body[i++] = 0xFF;
    //SPS数量为01，通过位运算得到0xE1
    packet->m_body[i++] = 0xE1;
    //SPS长度，占2个字节
    packet->m_body[i++] = (sps_len >> 8) & 0xff;
    packet->m_body[i++] = sps_len & 0xff;
    //SPS字节集
    memcpy(&packet->m_body[i], sps, sps_len);
    i += sps_len;
    //PPS数量
    packet->m_body[i++] = 0x01;
    //PPS长度，占两个字节
    packet->m_body[i++] = (pps_len >> 8) & 0xff;
    packet->m_body[i++] = (pps_len) & 0xff;
    //PPS字节集
    memcpy(&packet->m_body[i], pps, pps_len);
    //RTMPPacket包类型
    packet->m_packetType = RTMP_PACKET_TYPE_VIDEO;
    //RTMPPacket头类型
    packet->m_headerType = RTMP_PACKET_SIZE_MEDIUM;
    //Body大小
    packet->m_nBodySize = bodySize;
    //通道号
    packet->m_nChannel = 0x10;
    //时间戳
    packet->m_nTimeStamp = 0;
    //使用相对时间
    packet->m_hasAbsTimestamp = 0;
    //H264裸数据写入文件
    fs.write(packet->m_body, packet->m_nBodySize);
    return packet;
}

//创建存储H264视频裸数据的Packet
RTMPPacket *VideoChannel::createRawDataPacket(int type, uint8_t *payload, int payloadSize) {
    //去除H264的固定头，00 00 00 01或00 00 01
    if (payload[2] == 0x00) {
        payloadSize -= 4;
        payload += 4;
    } else {
        payloadSize -= 3;
        payload += 3;
    }
    //创建RTMPPacket
    int bodySize = 9 + payloadSize;
    RTMPPacket *packet = new RTMPPacket;
    RTMPPacket_Alloc(packet, bodySize);
    //01表示完整帧（IDR），02表示切片帧（SLICE），07表示x264编码，通过位运算得到0x17或0x27
    if (type == NAL_SLICE_IDR)
        packet->m_body[0] = 0x17;
    else
        packet->m_body[0] = 0x27;
    //00表示AVC Header，01表示AVC NALU
    packet->m_body[1] = 0x01;
    //时间戳
    packet->m_body[2] = 0x00;
    packet->m_body[3] = 0x00;
    packet->m_body[4] = 0x00;
    //NALU长度，占4个字节
    packet->m_body[5] = (payloadSize >> 24) & 0xff;
    packet->m_body[6] = (payloadSize >> 16) & 0xff;
    packet->m_body[7] = (payloadSize >> 8) & 0xff;
    packet->m_body[8] = (payloadSize) & 0xff;
    //NALU字节集
    memcpy(&packet->m_body[9], payload, payloadSize);
    //RTMPPacket包类型
    packet->m_packetType = RTMP_PACKET_TYPE_VIDEO;
    //RTMPPacket头类型
    packet->m_headerType = RTMP_PACKET_SIZE_LARGE;
    //Body大小
    packet->m_nBodySize = bodySize;
    //通道号
    packet->m_nChannel = 0x10;
    //使用相对时间
    packet->m_hasAbsTimestamp = 0;
    //相对时间戳
    packet->m_nTimeStamp = RTMP_GetTime() - startTime;
    return packet;
}
